# 导入所需的模块和类
from langchain.embeddings import CacheBackedEmbeddings
from langchain.storage import LocalFileStore
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain.embeddings.dashscope import DashScopeEmbeddings

from langchain_text_splitters import CharacterTextSplitter


class DB:
    def __init__(self):
        """
        初始化DB类，用于处理向量化和向量检索。
        """
        # 实例化 DashScope 嵌入器
        self.embeddings = DashScopeEmbeddings()

        # 初始化本地缓存存储器，缓存嵌入的结果，提升性能
        self.store = LocalFileStore("./cache/")

        # 创建支持缓存的嵌入器，避免重复计算
        self.cached_embedder = CacheBackedEmbeddings.from_bytes_store(
            self.embeddings, self.store, namespace=self.embeddings.model
        )
        print("嵌入器初始化成功：", self.cached_embedder)

    def add(self, chunks, key):
        """
        添加文档数据并向量化存储。
        :param chunks: 文档块（已分割的文档列表）
        :param key: 存储索引的名称
        """
        # 创建向量存储
        db = FAISS.from_documents(chunks, self.cached_embedder)
        # 保存向量存储到本地
        db.save_local(key)
        print(f"数据成功添加到索引：{key}")

    def search(self, ask, key, count):
        """
        搜索问题的相关文档。
        :param ask: 用户查询问题
        :param key: 使用的向量存储索引
        :param count: 返回的结果数量
        :return: 检索到的相关文档
        """
        # 加载本地向量存储
        db = FAISS.load_local(key, self.cached_embedder, allow_dangerous_deserialization=True)
        # 执行相似度搜索
        res = db.similarity_search(ask, k=count)
        return res


# 实例化 DB 类
myfaiss = DB()
